"""
XOR FEC implementation.
"""

from typing import List, Optional, Dict, Any, TYPE_CHECKING
import time
from ..logger_config import log_info, log_debug, log_warning, log_error, log_success, log_failure, log_progress, log_connection, log_media, log_stats

try:
    from .base import BaseFecCodec, FecConfig, FecPacket
except ImportError:
    from base import BaseFecCodec, FecConfig, FecPacket
if TYPE_CHECKING:
    try:
        from ..rtp import RtpPacket
    except ImportError:
        from rtp import RtpPacket


class XorFec(BaseFecCodec):
    """XOR FEC实现"""
    
    def __init__(self, config: FecConfig):
        super().__init__(config)
    
    def encode(self, packets: List['RtpPacket']) -> List[FecPacket]:
        """XOR编码"""
        if len(packets) != self.config.k:
            raise ValueError(f"Expected {self.config.k} packets, got {len(packets)}")
        
        # 提取RTP包负载
        payloads = [pkt.payload for pkt in packets]
        
        # XOR编码
        repair_payloads = self._xor_encode(payloads)
        
        # 生成FEC包
        fec_packets = []
        for i, repair_payload in enumerate(repair_payloads):
            fec_packet = FecPacket(
                fec_payload_id=self.sequence_number + i,
                repair_payload=repair_payload,
                protected_sequence_numbers=[pkt.sequence_number for pkt in packets],
                fec_type='xor'
            )
            fec_packets.append(fec_packet)
        
        self.sequence_number += len(fec_packets)
        self.stats['encoded_packets'] += len(fec_packets)
        
        return fec_packets
    
    def decode(self, original_packets: List[Optional['RtpPacket']], 
               fec_packets: List[FecPacket]) -> List[Optional['RtpPacket']]:
        """XOR解码"""
        self.stats['decoded_packets'] += 1
        
        if not fec_packets:
            return original_packets
        
        # XOR解码
        try:
            repaired_payloads = self._xor_decode(original_packets, fec_packets)
            
            # 创建修复的包
            result = []
            for i, payload in enumerate(repaired_payloads):
                if payload is not None:
                    # 创建修复的RTP包
                    repaired_packet = RtpPacket(
                        payload_type=96,
                        sequence_number=i,
                        timestamp=int(time.time() * 90000),
                        payload=payload
                    )
                    result.append(repaired_packet)
                    self.stats['recovered_packets'] += 1
                else:
                    result.append(None)
                    if original_packets[i] is None:
                        self.stats['failed_recoveries'] += 1
            
            return result
        except Exception:
            # 解码失败
            self.stats['failed_recoveries'] += sum(1 for p in original_packets if p is None)
            return original_packets
    
    def can_repair(self, lost_packets: List[int], 
                   available_fec_packets: List[FecPacket]) -> bool:
        """检查是否可以用可用FEC包修复丢失的包"""
        if not available_fec_packets:
            return False
        
        # XOR只能修复一个丢失的包
        return len(lost_packets) == 1
    
    def _xor_encode(self, payloads: List[bytes]) -> List[bytes]:
        """XOR编码"""
        if not payloads:
            return []
        
        # 找到最大长度
        max_length = max(len(p) for p in payloads)
        
        # 创建XOR结果
        xor_result = bytearray(max_length)
        for payload in payloads:
            for i, byte in enumerate(payload):
                xor_result[i] ^= byte
        
        return [bytes(xor_result)]
    
    def _xor_decode(self, original_packets: List[Optional['RtpPacket']], 
                   fec_packets: List[FecPacket]) -> List[Optional[bytes]]:
        """XOR解码"""
        if not fec_packets:
            return [p.payload if p else None for p in original_packets]
        
        # 找到丢失的包
        missing_indices = [i for i, p in enumerate(original_packets) if p is None]
        
        if len(missing_indices) != 1:
            # XOR只能修复一个丢失的包
            return [p.payload if p else None for p in original_packets]
        
        # 使用XOR修复
        missing_idx = missing_indices[0]
        fec_payload = fec_packets[0].repair_payload
        
        # 计算丢失的包
        result = []
        for i, packet in enumerate(original_packets):
            if i == missing_idx:
                # 修复丢失的包
                if packet is None:
                    # 计算丢失的包内容
                    xor_result = bytearray(fec_payload)
                    for j, other_packet in enumerate(original_packets):
                        if j != i and other_packet is not None:
                            for k, byte in enumerate(other_packet.payload):
                                if k < len(xor_result):
                                    xor_result[k] ^= byte
                    result.append(bytes(xor_result))
                else:
                    result.append(packet.payload)
            else:
                result.append(packet.payload if packet else None)
        
        return result